
#include "SysLib.h"

#ifndef OS_WIN

#include "Protocal.h"
#include "TunDev.h"
#include "Common.h"
#include "Config.h"
#include "SocketUtil.h"
#include "Firewall.h"

bool is_client;

PVpnConfig config = NULL;
PUserInfo *usersConf = NULL;
PTunDev tundev = NULL;
unsigned char network_mask_int[3];

// Attributes for server mode
int server_socket;

PThread route_thread;
PBuffer local_read_buffer;

int last_auth_seed = 0;
int connect_retry = 0;
long last_connect_time;
long cold_reconnect_time = -1;

void do_exit() {
    log_info("do_exit()");
    exit(1);
}


void * add_routes(void *args) {
    log_info("Adding routes");
    PRoute *push_routes = read_routes_cache(config->route_cache_file);
    int count = 0;
    if (push_routes) {
        PRoute *route;
        for (route=push_routes; *route; route++) {
            tundev_route(tundev, route[0]->address, route[0]->prefix);
            count++;
            route_delete(*route);
        }
        free(push_routes);
    }
    if (config->routes) {
        PRoute *route;
        for (route=config->routes; *route; route++) {
            tundev_route(tundev, route[0]->address, route[0]->prefix);
            count++;
        }
    }
    log_info("Add %d routes", count);
    return NULL;
}

void update_state(const char *state) {
    if (config->state_file != NULL) {
        char tmp[100];
        sprintf(tmp, "echo \"%s\" > %s", state, config->state_file);
        system(tmp);
    }
}

bool process_auth (PTunnel tunnel) {
    if (is_client) {
        client_auth(tunnel, config);
    } else {
        server_auth(tunnel, config);
    }

    bool init_client = false;
    if (tunnel->status == TSTATUS_ESTABLISH) {
        log_info("Auth pass, user: %s, client_ip: %s, server_ip: %s", tunnel->username, tunnel->client_ip, tunnel->server_ip);
        if (is_client) {
            ip_str_to_num(tunnel->client_ip, 4, tunnel->client_ip_num);
            bool hot_reconnect = tundev->local_ip && (strcmp(tundev->local_ip, tunnel->client_ip) == 0);
            if (hot_reconnect) {
                log_debug("Hot reconnect success");
            } else {
                if (tundev->local_ip) {
                    log_error("Ip change, exit");
                    do_exit();
                }

                if (tundev_ifconfig(tundev, tunnel->client_ip, tunnel->server_ip)) {
                    log_error("ifconfig error");
                    do_exit();
                }
                init_client = true;
                route_thread = thread_start(add_routes, NULL);
            }
            last_auth_seed = tunnel->auth_seed;
            update_state("Connected");
            cold_reconnect_time = 0;
        }
        if (init_client || !is_client) {
            int i;
            for (i=0; usersConf[i] != NULL; i++) {
                if (strcmp(usersConf[i]->username, tunnel->username) == 0) {
                    const char *script = usersConf[i]->script;
                    if (script && strlen(script)) {
                        char script_cmd[1024];
                        sprintf(script_cmd, "%s %s %s %s", script, tundev->dev,
                                tunnel->client_ip, tunnel->server_ip);
                        log_info("Execute up script: %s", script_cmd);
                        system(script_cmd);
                    }
                    break;
                }
            }
        }
    }
    return true;
}

bool check_network_mask(unsigned char *raw) {
    int i=0;
    for (i=0; i<3; i++) {
        if (network_mask_int[i] != (unsigned short)raw[i]) {
            return false;
        }
    }
    return true;
}

void do_init(int argc, char *argv[]) {
    tunnel_module_init();
    set_log_file(NULL);
    /**/
    if (argc <= 1) {
        printf("Run with myvpn <config_file>\n");
        do_exit();
    }
    config = read_config(argv[1]); /**/
    //config = read_config("d:\\myvpn.txt");
    if (config == NULL || !vpn_config_verify(config)) {
        log_error("Error verify config, exit");
        do_exit();
    }

    usersConf = config->users;
    set_log_level(config->log_level);
    set_log_file(config->log_file);

    log_info("Finish parsing parameters");

    tundev = tundev_alloc(config->dev);
    if (tundev == NULL) {
        log_error("Allocate TunDev error");
        do_exit();
    }

    local_read_buffer = buffer_new(config->read_buffer_size);
    if ((is_client = (config->mode == CLIENT))) {
    } else {
        ip_str_to_num(config->network, 3, network_mask_int);
        int optval = 1;
        if ( (server_socket = socket(AF_INET, SOCK_STREAM, 0)) < 0) {
            log_error("socket()");
            do_exit();
        }

        /* avoid EADDRINUSE error on bind() */
        if(setsockopt(server_socket, SOL_SOCKET, SO_REUSEADDR, (char *)&optval, sizeof(optval)) < 0) {
            log_error("setsockopt()");
            do_exit();
        }
        struct sockaddr_in local;
        memset(&local, 0, sizeof(local));
        local.sin_family = AF_INET;
        local.sin_addr.s_addr = htonl(INADDR_ANY);
        local.sin_port = htons(config->port);
        if (bind(server_socket, (struct sockaddr*) &local, sizeof(local)) < 0) {
            log_error("bind()");
            do_exit();
        }

        if (listen(server_socket, 5) < 0) {
            log_error("listen()");
            do_exit();
        }

        char local_ip[20];
        char remote_ip[20];
        sprintf(local_ip, "%s%d", config->network, 1);
        sprintf(remote_ip, "%s%d", config->network, 0);
        if (tundev_ifconfig(tundev, local_ip, remote_ip)) {
            log_error("TunDev ifconfig error");
            do_exit();
        }
        if (tundev_route(tundev, remote_ip, 24)) {
            log_error("Add route error");
            do_exit();
        }
    }
}

void check_client_connection() {
    PTunnel header = tunnel_get_chain_head();
    if (!header->next) {
        if (cold_reconnect_time > 0 && time(0) - cold_reconnect_time > 30) {
            do_exit();
        }
        if (cold_reconnect_time == 0) {
            cold_reconnect_time = time(0);
        }
        update_state("Connecting");
        if (time(0) - last_connect_time < 30) {
            sleep(min(connect_retry, 30));
            connect_retry ++;
        } else {
            connect_retry = 0;
        }
        last_connect_time = time(0);

        log_info("Create new client connection");
        // 如果没有tunnel则创建一个，并初始化连接状态
        PTunnel tunnel = tunnel_new(config->read_buffer_size,
                config->write_buffer_size);
        tunnel->username = string_alloc(config->users[0]->username);
        tunnel->client_auth_seed = last_auth_seed;
        if (!tunnel_connect(tunnel, config->host, config->port)) {
            log_error("Error connecting server");
            tunnel_delete(tunnel);
            return;
        }
    } else {
        header = header->next;
        if (!header->next && header->status == TSTATUS_ESTABLISH
                && time(0) - header->connect_time > config->renew_interval) {
            update_state("Renewing");
            log_info("Try renewing connection");
            PTunnel tunnel = tunnel_new(config->read_buffer_size,
                    config->write_buffer_size);
            tunnel->username = string_alloc(config->users[0]->username);
            tunnel->client_auth_seed = last_auth_seed;
            if (!tunnel_connect(tunnel, config->host, config->port)) {
                log_error("Error connecting server");
                tunnel_delete(tunnel);
                return;
            }
        }
    }
}

int main(int argc, char *argv[]) {
    do_init(argc, argv);

    int rwcount;
    struct timeval timeout;
    long now;
    fd_set rd_set;
    fd_set wt_set;
    char tun_buffer[1600];
    int max_fd;

    while(1) {
        if (is_client) check_client_connection();

        FD_ZERO(&rd_set);
        FD_ZERO(&wt_set);

        int rcount = 1;
        int wcount = 0;

        FD_SET(tundev->fd, &rd_set);
        int max_fd = tundev->fd;

        if (!is_client) {
            FD_SET(server_socket, &rd_set);
            max_fd = max(max_fd, server_socket);
            rcount += 1;
        }

        int i;

        PTunnel iter = NULL;

        iter = tunnel_get_chain_head();
        while (iter->next) {
            PTunnel tunnel = iter->next;
            iter = iter->next;

            if (tunnel_is_readable_status(tunnel)) {
                FD_SET(tunnel->socket->fd, &rd_set);
                max_fd = max(max_fd, tunnel->socket->fd);
                rcount ++;
            }

            if (tunnel_is_writeable_status(tunnel)) {
                FD_SET(tunnel->socket->fd, &wt_set);
                max_fd = max(max_fd, tunnel->socket->fd);
                wcount ++;
            }
        }

        timeout.tv_sec = 5;
        timeout.tv_usec = 0;
        rwcount = select(max_fd + 1, &rd_set, &wt_set, NULL, &timeout);
        if (rwcount < 0) {
            if (errno != EINTR) {
                log_error("select return %d, errno: %d", rwcount, errno);
                do_exit();
            }
            continue;
        }
        log_trace("Select r: %d, w: %d, ret: %d", rcount, wcount, rwcount);
        now = time(0);
        if (is_client == false && FD_ISSET(server_socket, &rd_set)) {
            PTunnel tunnel = tunnel_accept(server_socket, config->write_buffer_size);
        }
        if (FD_ISSET(tundev->fd, &rd_set)) {
            rwcount = (int)read(tundev->fd, tun_buffer, 1500);
            if (rwcount < 0) {
                log_error("Read from tun return %d, %d", rwcount, errno);
                do_exit();
            }
            log_trace("Read %d bytes from tun", rwcount);
            PTunnel tunnel = NULL;
            if (is_client) {
                unsigned int dst = ((unsigned char *)tun_buffer)[15];
                tunnel = tunnel_get_by_ip_id(dst);
                if (tunnel && !filter_upstream_pack((unsigned char*)tun_buffer, rwcount, tunnel)) {
                    tunnel = NULL;
                }
            } else {
                bool valid = true;
                unsigned int dst = ((unsigned char *)tun_buffer)[19];
                if (!check_network_mask((unsigned char *)(tun_buffer + 16))) {
                    valid = false;
                }
                tunnel = tunnel_get_by_ip_id(dst);
                if (!(valid && tunnel && tunnel->status == TSTATUS_ESTABLISH)) {
                    valid = false;
                }
                if (!valid) {
                    tunnel = NULL;
                    log_error("Drop ip packet for no valid tunnel, dst: %u.%u.%u.%u",
                            ((unsigned char *)tun_buffer)[16],
                            ((unsigned char *)tun_buffer)[17],
                            ((unsigned char *)tun_buffer)[18],
                            dst);
                }
            }

            if (!tunnel) {
                continue;
            }

            if (0 > encode_data(tun_buffer, rwcount, tunnel->write_buffer)) {
                log_error("Error handling receive data");
                tunnel_close(tunnel, true);
                continue;
            }
        }
        iter = tunnel_get_chain_head();
        while (iter->next) {
            PTunnel tunnel = iter->next;

            PSocketWrapper sock = tunnel->socket;
            int net_fd = sock->fd;
            Buffer *read_buffer = tunnel->read_buffer;
            Buffer *write_buffer = tunnel->write_buffer;

            if(tunnel_is_readable_status(tunnel) && FD_ISSET(net_fd, &rd_set)) {
                buffer_clear(local_read_buffer);
                buffer_copy(local_read_buffer, read_buffer);
                rwcount = socket_read(sock, local_read_buffer);

                if (rwcount <= 0) {
                    log_error("Read from net return %d, errno: %d", rwcount, errno);
                    tunnel_close(tunnel, true);
                    continue;
                }
                if (tunnel->status != TSTATUS_ESTABLISH) {
                    log_debug("Read %d bytes from network, clientIp: %s, bufFrom: %d, bufTo: %d, remain: %d", rwcount, tunnel->client_ip, local_read_buffer->from, local_read_buffer->to, buffer_writable(local_read_buffer));
                } else {
                    log_trace("Read %d bytes from network, clientIp: %s, bufFrom: %d, bufTo: %d, remain: %d", rwcount, tunnel->client_ip, local_read_buffer->from, local_read_buffer->to, buffer_writable(local_read_buffer));
                }
                tunnel->last_from_net = now;

                if (tunnel->status >= TSTATUS_AUTH_BEGIN && tunnel->status < TSTATUS_ESTABLISH) {
                    tunnel->read_buffer = local_read_buffer;
                    if (!process_auth(tunnel)) {
                        tunnel->read_buffer = read_buffer;
                        tunnel_close(tunnel, false);
                        continue;
                    } else {
                        tunnel->read_buffer = read_buffer;
                    }
                } else {
                    int plen;
                    while ((plen = decode_data(local_read_buffer, tun_buffer)) >= 0) {
                        if (plen == PTYPE_HEARTBEAT) {
                            log_debug("Receive heartbeat");
                        } else {
                            if ((!is_client) && (!filter_upstream_pack((unsigned char *)tun_buffer, plen, tunnel))) {
                            } else {
                                int nwrite = tundev_write(tundev, tun_buffer, plen);
                                if (plen != nwrite) {
                                    log_error("Write to tun len not match, exp: %d, ret: %d", plen, nwrite);
                                    do_exit();
                                }
                            }
                        }
                    }
                    if (plen <= PTYPE_ILLEGAL) {
                        if (plen == PTYPE_END) {
                            log_debug("Connection close gracefully");
                        } else {
                            log_error("Error decoding receive data, error: %d", plen);
                        }
                        tunnel_close(tunnel, true);
                        continue;
                    }
                }
                buffer_copy(read_buffer, local_read_buffer);
            }

            if (tunnel_is_writeable_status(tunnel) && FD_ISSET(net_fd, &wt_set)) {
                if (tunnel->status == TSTATUS_CONNECTING) {
                    if (tunnel_check_nio_connect(tunnel)) {
                        tunnel->status = TSTATUS_AUTH_BEGIN;
                        if (!process_auth(tunnel)) {
                            tunnel_close(tunnel, true);
                        }
                    } else {
                        tunnel_close(tunnel, true);
                    }
                    continue;
                }

                rwcount = socket_write(sock, write_buffer);
                if (rwcount <= 0) {
                    log_error("Write to net return %d, errno: %d", rwcount, errno);
                    tunnel_close(tunnel, true);
                    continue;
                }
                if (tunnel->status != TSTATUS_ESTABLISH) {
                    log_debug("Send %d bytes to network, clientIp: %s, bufFrom: %d, bufTo: %d, remain: %d", rwcount, tunnel->client_ip, write_buffer->from, write_buffer->to, buffer_readable(write_buffer));
                } else {
                    log_trace("Send %d bytes to network, clientIp: %s, bufFrom: %d, bufTo: %d, remain: %d", rwcount, tunnel->client_ip, write_buffer->from, write_buffer->to, buffer_readable(write_buffer));
                }
                tunnel->last_to_net = now;
            }

            if (tunnel->status == TSTATUS_ESTABLISH) {
                if (now - tunnel->connect_time > 3600) {
                    log_info("Close old connection: %s", tunnel->username);
                    tunnel_close(tunnel, false);
                    continue;
                }

                if (now - tunnel->last_to_net > tunnel->heartbeat) {
                    log_debug("Trigger heartbeat");
                    encode_data(tun_buffer, PTYPE_HEARTBEAT, write_buffer);
                }
            }
            iter = iter->next;
        }

        iter = tunnel_get_chain_head();
        while (iter->next) {
            PTunnel tunnel = iter->next;
            bool delete = false;
            if (tunnel->status == TSTATUS_CLOSING
                    && 0 == buffer_readable(tunnel->write_buffer)) {
                delete = true;
            } else if (now - tunnel->last_to_net > tunnel->timeout) {
                log_error("Send timeout, user: %s", tunnel->username);
                delete = true;
            } else if (now - tunnel->last_from_net > tunnel->timeout) {
                log_error("Heartbeat timeout, user: %s", tunnel->username);
                delete = true;
            }
            if (delete) {
                tunnel_delete(tunnel);
            } else {
                iter = iter->next;
            }
        }
    }

    return(0);
}

#endif
